# _*_ coding: utf-8 _*_
# Codes from Baidu AI, Inc.
# Licensed under the Apache License, Version 2.0 (the "License")
# basic utils for helixwrapper
# Author: MiqroEra Shibo

import numpy as np
import os
import random
import json

from pgl.utils.data import Dataloader


def mp_pool_map(list_input, func, num_workers):
    """
    The function is to generate a list of results from a list of inputs
    using multiprocessing pool.
    params list_input: list of inputs
    params func: function to be applied to each input
    params num_workers: number of workers
    return: list of results
    """

    class _CollateFn(object):
        def __init__(self, func):
            self.func = func

        def __call__(self, data_list):
            new_data_list = []
            for data in data_list:
                index, input = data
                new_data_list.append((index, self.func(input)))
            return new_data_list

    list_new_input = [(index, x) for index, x in enumerate(list_input)]
    data_gen = Dataloader(list_new_input, batch_size=8, shuffle=False, num_workers=num_workers,
                          collate_fn=_CollateFn(func))

    list_output = []
    for sub_output in data_gen:
        list_output += sub_output
    list_output = sorted(list_output, key=lambda x: x[0])
    list_output = [x[1] for x in list_output]
    return list_output


def load_json_config(path):
    """
    The function is to load json config file.
    params path: path of json config file
    return: config
    """
    return json.load(open(path, 'r'))
